import torch
import torchvision.models as models
from torch.optim import Adam
from memory_profiler import FakeTensorMemoryProfilerMode, tensor_storage_id
from torch._subclasses import FakeTensorMode

if __name__=="__main__":
    MB = 2 ** 20
    GB = 2 ** 30

    MEMORY_LIMIT = 16 * GB



    def func(batch_size):
        print(f"Running batch size {batch_size}")
        with FakeTensorMode(allow_non_fake_inputs=True):
            with FakeTensorMemoryProfilerMode() as ftmp:
                # 初始化 ResNet50 模型
                model = models.resnet50()
                optimizer = Adam(model.parameters(), lr=1e-5)
                ftmp.add_marker("model_init_boundary")
            
                # 生成随机图像数据 (batch_size, 3, 224, 224)
                input_tensor = torch.rand(batch_size, 3, 224, 224, requires_grad=True)
                allocated_memory = torch.cuda.memory_allocated()
                print(f"Allocated memory: {allocated_memory / 1024 ** 2:.5f} MB")
                for i in range(3):
                    output = model(input_tensor)
                    print(f"GB after forward: {ftmp.max_memory / GB}")
                    ftmp.add_marker(f"fw_bw_boundary_{i}")
                    output.sum().backward()
                    ftmp.add_marker(f"bw_step_boundary_{i}")
                    print(f"GB after backward: {ftmp.max_memory / GB}")
                    optimizer.step()
                    ftmp.add_marker(f"step_boundary_{i}")
                    print(f"GB after step: {ftmp.max_memory / GB}")

                ftmp.draw_varies()
                return ftmp.max_memory

    with torch.device("cuda:0"):
        func(2)
        allocated_memory = torch.cuda.memory_allocated()
        print(f"Allocated memory: {allocated_memory / 1024 ** 2:.5f} MB")
